import argparse
import os
import re

import pandas

parser = argparse.ArgumentParser()
parser.add_argument("file", type=str, help="path to the raw data file")
parser.add_argument(
    "--column",
    type=str,
    help="column names",
    choices=["location", "department", "type"],
)
args = parser.parse_args()

output_dir = os.path.join("data", args.column)
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

columns = ["input_1", "input_2", "location", "department", "type"]
raw_dataframe = pandas.read_excel(args.file, names=columns)

dataframe = raw_dataframe[["input_1", "input_2", args.column]].copy()
# add a new column "input" which is the concatenation of input_1 and input_2

dataframe["input"] = dataframe["input_1"] + " " + dataframe["input_2"]
dataframe = dataframe.drop(columns=["input_1", "input_2"])
# rename the column to "label"
dataframe = dataframe.rename(columns={args.column: "label"})
# change the order of columns
dataframe = dataframe[["input", "label"]]

dataframe = dataframe.sample(frac=1).reset_index(drop=True)
# filter out empty input and label
dataframe = dataframe[dataframe["input"].apply(lambda x: len(x) > 0)]
dataframe = dataframe[dataframe["label"].apply(lambda x: len(x) > 0)]

# replace \s+ with space in input
dataframe["input"] = dataframe["input"].apply(lambda x: re.sub(r"\s+", " ", x))

total = dataframe.shape[0]
test_size = min(50, total // 10)

dev = dataframe.iloc[:test_size]
test = dataframe.iloc[test_size : 2 * test_size]
train = dataframe.iloc[2 * test_size :]

# to tsv
train.to_csv(os.path.join(output_dir, "train.txt"), sep="\t", index=False, header=False)
dev.to_csv(os.path.join(output_dir, "dev.txt"), sep="\t", index=False, header=False)
test.to_csv(os.path.join(output_dir, "test.txt"), sep="\t", index=False, header=False)

# all_labels
all_labels = dataframe["label"].unique()
with open(os.path.join(output_dir, "label.txt"), "w") as f:
    for label in all_labels:
        print(label, file=f)

# data.txt
with open(os.path.join(output_dir, "data.txt"), "w") as f:
    for row in test.itertuples():
        print(row[1], file=f)
